//===--- ClosureLifetimeFixup.cpp - Fixup the lifetime of closures --------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "closure-lifetime-fixup"

#include "polarphp/pil/lang/DebugUtils.h"
#include "polarphp/pil/lang/InstructionUtils.h"
#include "polarphp/pil/lang/PILArgument.h"
#include "polarphp/pil/lang/PILBuilder.h"
#include "polarphp/pil/lang/PILInstruction.h"
#include "polarphp/pil/optimizer/passmgr/Passes.h"
#include "polarphp/pil/optimizer/passmgr/Transforms.h"
#include "polarphp/pil/optimizer/utils/CFGOptUtils.h"
#include "polarphp/pil/optimizer/utils/InstOptUtils.h"
#include "polarphp/pil/optimizer/utils/PILSSAUpdater.h"
#include "polarphp/pil/optimizer/utils/StackNesting.h"

#include "llvm/Support/CommandLine.h"

llvm::cl::opt<bool> DisableConvertEscapeToNoEscapeSwitchEnumPeephole(
   "pil-disable-convert-escape-to-noescape-switch-peephole",
   llvm::cl::init(false),
   llvm::cl::desc(
      "Disable the convert_escape_to_noescape switch enum peephole. "),
   llvm::cl::Hidden);

using namespace polar;

/// Given an optional diamond, return the bottom of the diamond.
///
/// That is given that sei is in bb0,
///
///   /---> bb1 ---\
///  /              \
/// bb0              ---> bb3
///  \              /
///   \---> bb2 ---/
///
/// this routine will return bb3.
static PILBasicBlock *getOptionalDiamondSuccessor(SwitchEnumInst *sei) {
   auto numSuccs = sei->getNumSuccessors();
   if (numSuccs != 2)
      return nullptr;
   auto *succSome = sei->getCase(0).second;
   auto *succNone = sei->getCase(1).second;
   if (succSome->args_size() != 1)
      std::swap(succSome, succNone);

   if (succSome->args_size() != 1 || succNone->args_size() != 0)
      return nullptr;

   auto *succ = succSome->getSingleSuccessorBlock();
   if (!succ)
      return nullptr;

   if (succNone == succ)
      return succ;

   succNone = succNone->getSingleSuccessorBlock();
   if (succNone == succ)
      return succ;

   if (succNone == nullptr)
      return nullptr;

   succNone = succNone->getSingleSuccessorBlock();
   if (succNone == succ)
      return succ;

   return nullptr;
}

/// Find a safe insertion point for closure destruction. We might create a
/// closure that captures self in deinit of self. In this situation it is not
/// safe to destroy the closure after we called super deinit. We have to place
/// the closure destruction before that call.
///
///  %deinit = objc_super_method %0 : $C, #A.deinit!deallocator.foreign
///  %super = upcast %0 : $C to $A
///  apply %deinit(%super) : $@convention(objc_method) (A) -> ()
///  end_lifetime %super : $A
static PILInstruction *getDeinitSafeClosureDestructionPoint(PILBasicBlock *bb) {
   for (auto &i : llvm::reverse(*bb)) {
      if (auto *endLifetime = dyn_cast<EndLifetimeInst>(&i)) {
         auto *superInstance = endLifetime->getOperand()->getDefiningInstruction();
         assert(superInstance && "Expected an instruction");
         return superInstance;
      }
   }
   return bb->getTerminator();
}

static void findReachableExitBlocks(PILInstruction *i,
                                    SmallVectorImpl<PILBasicBlock *> &result) {
   SmallVector<PILBasicBlock *, 32> worklist;
   SmallPtrSet<PILBasicBlock *, 8> visitedBlocks;

   visitedBlocks.insert(i->getParent());
   worklist.push_back(i->getParent());

   while (!worklist.empty()) {
      auto *bb = worklist.pop_back_val();
      if (bb->getTerminator()->isFunctionExiting()) {
         result.push_back(bb);
         continue;
      }
      llvm::copy_if(bb->getSuccessorBlocks(), std::back_inserter(worklist),
                    [&](PILBasicBlock *bb) {
                       return visitedBlocks.insert(bb).second;
                    });
   }
}

/// We use this to ensure that we properly handle recursive cases by revisiting
/// phi nodes that we are tracking. This just makes it easier to reproduce in a
/// test case.
static llvm::cl::opt<bool> ReverseInitialWorklist(
   "sil-closure-lifetime-fixup-reverse-phi-order", llvm::cl::init(false),
   llvm::cl::desc(
      "Reverse the order in which we visit phis for testing purposes"),
   llvm::cl::Hidden);

// Finally, we need to prune phis inserted by the SSA updater that
// only take the .none from the entry block. This means that they are
// not actually reachable from the .some() so we know that we do not
// need to lifetime extend there at all. As an additional benefit, we
// eliminate the need to balance these arguments to satisfy the
// ownership verifier. This occurs since arguments are a place in PIL
// where the trivialness of an enums case is erased.
static void
cleanupDeadTrivialPhiArgs(PILValue initialValue,
                          SmallVectorImpl<PILPhiArgument *> &insertedPhis) {
   // Just for testing purposes.
   if (ReverseInitialWorklist) {
      std::reverse(insertedPhis.begin(), insertedPhis.end());
   }
   SmallVector<PILPhiArgument *, 8> worklist(insertedPhis.begin(),
                                             insertedPhis.end());
   sortUnique(insertedPhis);
   SmallVector<PILValue, 8> incomingValues;

   while (!worklist.empty()) {
      // Clear the incoming values array after each iteration.
      POLAR_DEFER { incomingValues.clear(); };

      auto *phi = worklist.pop_back_val();
      {
         auto it = lower_bound(insertedPhis, phi);
         if (it == insertedPhis.end() || *it != phi)
            continue;
      }

      // TODO: When we split true phi arguments from transformational terminators,
      // this will always succeed and the assert can go away.
      bool foundPhiValues = phi->getIncomingPhiValues(incomingValues);
      (void)foundPhiValues;
      assert(foundPhiValues && "Should always have 'true' phi arguments since "
                               "these were inserted by the SSA updater.");
      if (llvm::any_of(incomingValues,
                       [&](PILValue v) { return v != initialValue; }))
         continue;

      // Remove it from our insertedPhis list to prevent us from re-visiting this.
      {
         auto it = lower_bound(insertedPhis, phi);
         assert((it != insertedPhis.end() && *it == phi) &&
                "Should have found the phi");
         insertedPhis.erase(it);
      }

      // See if any of our users are branch or cond_br. If so, we may have
      // exposed additional unneeded phis. Add it back to the worklist in such a
      // case.
      for (auto *op : phi->getUses()) {
         auto *user = op->getUser();

         if (!isa<BranchInst>(user) && !isa<CondBranchInst>(user))
            continue;

         auto *termInst = cast<TermInst>(user);
         for (auto succBlockArgList : termInst->getSuccessorBlockArguments()) {
            llvm::copy_if(succBlockArgList, std::back_inserter(worklist),
                          [&](PILPhiArgument *succArg) -> bool {
                             auto it = lower_bound(insertedPhis, succArg);
                             return it != insertedPhis.end() && *it == succArg;
                          });
         }
      }

      // Then RAUW the phi with the entryBlockOptionalNone and erase the
      // argument.
      phi->replaceAllUsesWith(initialValue);
      erasePhiArgument(phi->getParent(), phi->getIndex());
   }
}

/// Extend the lifetime of the convert_escape_to_noescape's operand to the end
/// of the function.
///
/// NOTE: Since we are lifetime extending a copy that we have introduced, we do
/// not need to consider destroy_value emitted by PILGen unlike
/// copy_block_without_escaping which consumes its sentinel parameter. Unlike
/// that case where we have to consider that destroy_value, we have a simpler
/// time here.
static void extendLifetimeToEndOfFunction(PILFunction &fn,
                                          ConvertEscapeToNoEscapeInst *cvt) {
   auto escapingClosure = cvt->getOperand();
   auto escapingClosureTy = escapingClosure->getType();
   auto optionalEscapingClosureTy = PILType::getOptionalType(escapingClosureTy);
   auto loc = RegularLocation::getAutoGeneratedLocation();

   // If our Cvt is in the initial block, we do not need to use the SSA updater
   // since we know Cvt can not be in a loop and must dominate all exits
   // (*). Just insert a copy of the escaping closure at the Cvt and destroys at
   // the exit blocks of the function.
   //
   // (*) In fact we can't use the PILSSAUpdater::GetValueInMiddleOfBlock.
   if (cvt->getParent() == cvt->getFunction()->getEntryBlock()) {
      auto *innerCVI =
         PILBuilderWithScope(cvt).createCopyValue(loc, escapingClosure);
      cvt->setLifetimeGuaranteed();
      cvt->setOperand(innerCVI);
      SmallVector<PILBasicBlock *, 4> exitingBlocks;
      fn.findExitingBlocks(exitingBlocks);

      for (auto *block : exitingBlocks) {
         auto *safeDestructPoint = getDeinitSafeClosureDestructionPoint(block);
         PILBuilderWithScope(safeDestructPoint).createDestroyValue(loc, innerCVI);
      }
      return;
   }

   // Ok. At this point we know that Cvt is not in the entry block... so we can
   // use PILSSAUpdater::GetValueInMiddleOfBlock() to extend the object's
   // lifetime respecting loops.
   SmallVector<PILPhiArgument *, 8> insertedPhis;
   PILSSAUpdater updater(&insertedPhis);
   updater.Initialize(optionalEscapingClosureTy);

   // Create an Optional<() -> ()>.none in the entry block of the function and
   // add it as an available value to the SSAUpdater.
   //
   // Since we know that Cvt is not in the entry block and this must be, we know
   // that it is safe to use the SSAUpdater's getValueInMiddleOfBlock with this
   // value.
   PILValue entryBlockOptionalNone = [&]() -> PILValue {
      PILBuilderWithScope b(fn.getEntryBlock()->begin());
      return b.createOptionalNone(loc, optionalEscapingClosureTy);
   }();
   updater.AddAvailableValue(fn.getEntryBlock(), entryBlockOptionalNone);

   // Create a copy of the convert_escape_to_no_escape and add it as an available
   // value to the SSA updater.
   //
   // NOTE: The SSAUpdater does not support providing multiple values in the same
   // block without extra work. So the fact that Cvt is not in the entry block
   // means that we don't have to worry about overwriting the .none value.
   auto *cvi = [&]() -> CopyValueInst * {
      auto *innerCVI =
         PILBuilderWithScope(cvt).createCopyValue(loc, escapingClosure);
      cvt->setLifetimeGuaranteed();
      cvt->setOperand(innerCVI);
      PILBuilderWithScope b(std::next(cvt->getIterator()));
      updater.AddAvailableValue(
         cvt->getParent(),
         b.createOptionalSome(loc, innerCVI, optionalEscapingClosureTy));
      return innerCVI;
   }();

   // Then we use the SSA updater to insert a destroy_value before the cvt and at
   // the reachable exit blocks.
   SmallVector<PILBasicBlock *, 4> exitingBlocks;
   findReachableExitBlocks(cvt, exitingBlocks);

   {
      // Before the copy value, insert an extra destroy_value to handle
      // loops. Since we used our enum value this is safe.
      PILValue v = updater.GetValueInMiddleOfBlock(cvi->getParent());
      PILBuilderWithScope(cvi).createDestroyValue(loc, v);
   }

   for (auto *block : exitingBlocks) {
      auto *safeDestructionPt = getDeinitSafeClosureDestructionPoint(block);
      PILValue v = updater.GetValueAtEndOfBlock(block);
      PILBuilderWithScope(safeDestructionPt).createDestroyValue(loc, v);
   }

   // Finally, we need to prune phis inserted by the SSA updater that only take
   // the .none from the entry block.
   //
   // TODO: Should we sort inserted phis before or after we initialize
   // the worklist or maybe backwards? We should investigate how the
   // SSA updater adds phi nodes to this list to resolve this question.
   cleanupDeadTrivialPhiArgs(entryBlockOptionalNone, insertedPhis);
}

static PILInstruction *lookThroughRebastractionUsers(
   PILInstruction *inst,
   llvm::DenseMap<PILInstruction *, PILInstruction *> &memoized) {

   if (inst == nullptr)
      return nullptr;

   // Try a cached lookup.
   auto res = memoized.find(inst);
   if (res != memoized.end())
      return res->second;

   // Cache recursive results.
   auto memoizeResult = [&](PILInstruction *from, PILInstruction *toResult) {
      memoized[from] = toResult;
      return toResult;
   };

   // If we have a convert_function, just look at its user.
   if (auto *cvt = dyn_cast<ConvertFunctionInst>(inst))
      return memoizeResult(inst, lookThroughRebastractionUsers(
         getSingleNonDebugUser(cvt), memoized));
   if (auto *cvt = dyn_cast<ConvertEscapeToNoEscapeInst>(inst))
      return memoizeResult(inst, lookThroughRebastractionUsers(
         getSingleNonDebugUser(cvt), memoized));

   // If we have a partial_apply user look at its single (non release) user.
   auto *pa = dyn_cast<PartialApplyInst>(inst);
   if (!pa)
      return inst;

   PILInstruction *singleNonDebugNonRefCountUser = nullptr;
   for (auto *use : getNonDebugUses(pa)) {
      auto *user = use->getUser();
      if (onlyAffectsRefCount(user))
         continue;
      if (singleNonDebugNonRefCountUser) {
         singleNonDebugNonRefCountUser = nullptr;
         break;
      }
      singleNonDebugNonRefCountUser = user;
   }

   return memoizeResult(inst, lookThroughRebastractionUsers(
      singleNonDebugNonRefCountUser, memoized));
}

/// Insert a mark_dependence for any non-trivial argument of a partial_apply.
static PILValue insertMarkDependenceForCapturedArguments(PartialApplyInst *pai,
                                                         PILBuilder &b) {
   PILValue curr(pai);
   // Mark dependence on all non-trivial arguments.
   for (auto &arg : pai->getArgumentOperands()) {
      if (arg.get()->getType().isTrivial(*pai->getFunction()))
         continue;
      curr = b.createMarkDependence(pai->getLoc(), curr, arg.get());
   }

   return curr;
}

/// Rewrite a partial_apply convert_escape_to_noescape sequence with a single
/// apply/try_apply user to a partial_apply [stack] terminated with a
/// dealloc_stack placed after the apply.
///
///   %p = partial_apply %f(%a, %b)
///   %ne = convert_escape_to_noescape %p
///   apply %f2(%p)
///   destroy_value %p
///
///    =>
///
///   %p = partial_apply [stack] %f(%a, %b)
///   %md = mark_dependence %p on %a
///   %md2 = mark_dependence %md on %b
///   apply %f2(%md2)
///   dealloc_stack %p
///   destroy_value %a
///   destroy_value %b
///
/// Note: If the rewrite succeeded we have inserted a dealloc_stack. This
/// dealloc_stack still needs to be balanced with other dealloc_stacks i.e the
/// caller needs to use the StackNesting utility to update the dealloc_stack
/// nesting.
static bool tryRewriteToPartialApplyStack(
   PILLocation &loc, PartialApplyInst *origPA,
   ConvertEscapeToNoEscapeInst *cvt, PILInstruction *singleApplyUser,
   PILBasicBlock::iterator &advanceIfDelete,
   llvm::DenseMap<PILInstruction *, PILInstruction *> &memoized) {

   auto *convertOrPartialApply = cast<SingleValueInstruction>(origPA);
   if (cvt->getOperand() != origPA)
      convertOrPartialApply = cast<ConvertFunctionInst>(cvt->getOperand());

   // Whenever we delete an instruction advance the iterator and remove the
   // instruction from the memoized map.
   auto saveDeleteInst = [&](PILInstruction *i) {
      if (&*advanceIfDelete == i)
         advanceIfDelete++;
      memoized.erase(i);
      i->eraseFromParent();
   };

   // Look for a single non ref count user of the partial_apply.
   SmallVector<PILInstruction *, 8> refCountInsts;
   PILInstruction *singleNonDebugNonRefCountUser = nullptr;
   for (auto *use : getNonDebugUses(convertOrPartialApply)) {
      auto *user = use->getUser();
      if (onlyAffectsRefCount(user)) {
         refCountInsts.push_back(user);
         continue;
      }
      if (singleNonDebugNonRefCountUser)
         return false;
      singleNonDebugNonRefCountUser = user;
   }

   PILBuilderWithScope b(cvt);

   // The convert_escape_to_noescape is the only user of the partial_apply.
   // Convert to a partial_apply [stack].
   SmallVector<PILValue, 8> args;
   for (auto &arg : origPA->getArgumentOperands())
      args.push_back(arg.get());
   auto newPA = b.createPartialApply(
      origPA->getLoc(), origPA->getCallee(), origPA->getSubstitutionMap(), args,
      origPA->getType().getAs<PILFunctionType>()->getCalleeConvention(),
      PartialApplyInst::OnStackKind::OnStack);

   // Insert mark_dependence for any non-trivial operand to the partial_apply.
   auto closure = insertMarkDependenceForCapturedArguments(newPA, b);

   // Optionally, replace the convert_function instruction.
   if (auto *convert = dyn_cast<ConvertFunctionInst>(convertOrPartialApply)) {
      auto origTy = convert->getType().castTo<PILFunctionType>();
      auto origWithNoEscape = PILType::getPrimitiveObjectType(
         origTy->getWithExtInfo(origTy->getExtInfo().withNoEscape()));
      closure = b.createConvertFunction(convert->getLoc(), closure,
                                        origWithNoEscape, false);
      convert->replaceAllUsesWith(closure);
   }

   // Replace the convert_escape_to_noescape uses with the new
   // partial_apply [stack].
   cvt->replaceAllUsesWith(closure);
   saveDeleteInst(cvt);

   // Delete the ref count operations on the original partial_apply.
   for (auto *refInst : refCountInsts)
      saveDeleteInst(refInst);
   convertOrPartialApply->replaceAllUsesWith(newPA);
   if (convertOrPartialApply != origPA)
      saveDeleteInst(convertOrPartialApply);
   saveDeleteInst(origPA);

   // Insert destroys of arguments after the apply and the dealloc_stack.
   if (auto *apply = dyn_cast<ApplyInst>(singleApplyUser)) {
      auto insertPt = std::next(PILBasicBlock::iterator(apply));
      // Don't insert dealloc_stacks at unreachable.
      if (isa<UnreachableInst>(*insertPt))
         return true;
      PILBuilderWithScope b3(insertPt);
      b3.createDeallocStack(loc, newPA);
      insertDestroyOfCapturedArguments(newPA, b3);
   } else if (auto *tai = dyn_cast<TryApplyInst>(singleApplyUser)) {
      for (auto *succBB : tai->getSuccessorBlocks()) {
         PILBuilderWithScope b3(succBB->begin());
         b3.createDeallocStack(loc, newPA);
         insertDestroyOfCapturedArguments(newPA, b3);
      }
   } else {
      llvm_unreachable("Unknown FullApplySite instruction kind");
   }
   return true;
}

static PILValue skipConvert(PILValue v) {
   auto *cvt = dyn_cast<ConvertFunctionInst>(v);
   if (!cvt)
      return v;
   auto *pa = dyn_cast<PartialApplyInst>(cvt->getOperand());
   if (!pa || !pa->hasOneUse())
      return v;
   return pa;
}

static bool tryExtendLifetimeToLastUse(
   ConvertEscapeToNoEscapeInst *cvt,
   llvm::DenseMap<PILInstruction *, PILInstruction *> &memoized,
   PILBasicBlock::iterator &advanceIfDelete) {
   // If there is a single user that is an apply this is simple: extend the
   // lifetime of the operand until after the apply.
   auto *singleUser = lookThroughRebastractionUsers(cvt, memoized);
   if (!singleUser)
      return false;

   // Handle an apply.
   if (auto singleApplyUser = FullApplySite::isa(singleUser)) {
      // FIXME: Don't know how-to handle begin_apply/end_apply yet.
      if (isa<BeginApplyInst>(singleApplyUser.getInstruction())) {
         return false;
      }

      auto loc = RegularLocation::getAutoGeneratedLocation();
      auto origPA = dyn_cast<PartialApplyInst>(skipConvert(cvt->getOperand()));
      if (origPA && tryRewriteToPartialApplyStack(
         loc, origPA, cvt, singleApplyUser.getInstruction(),
         advanceIfDelete, memoized))
         return true;

      // Insert a copy at the convert_escape_to_noescape [not_guaranteed] and
      // change the instruction to the guaranteed form.
      auto escapingClosure = cvt->getOperand();
      auto *closureCopy =
         PILBuilderWithScope(cvt).createCopyValue(loc, escapingClosure);
      cvt->setLifetimeGuaranteed();
      cvt->setOperand(closureCopy);

      // Insert a destroy after the apply.
      if (auto *apply = dyn_cast<ApplyInst>(singleApplyUser.getInstruction())) {
         auto insertPt = std::next(PILBasicBlock::iterator(apply));
         PILBuilderWithScope(insertPt).createDestroyValue(loc, closureCopy);

      } else if (auto *tai =
         dyn_cast<TryApplyInst>(singleApplyUser.getInstruction())) {
         for (auto *succBB : tai->getSuccessorBlocks()) {
            PILBuilderWithScope(succBB->begin())
               .createDestroyValue(loc, closureCopy);
         }
      } else {
         llvm_unreachable("Unknown FullApplySite instruction kind");
      }
      return true;
   }
   return false;
}

/// Ensure the lifetime of the closure across a two step optional conversion
/// from:
///
///   optional<@escaping () -> ()>
///
/// to:
///
///   optional<@noescape () -> ()>
///
/// to:
///
///   optional<@noescape @convention(block) () -> ()>
///
/// and all uses of the block. The pattern that we are looking for is:
///
///                            switch_enum %optional_closure           (1)
///                           /           \
///   %trivial_closure = CVT %closure     nil                          (2)
///                           \           /
///                          switch_enum %optional_trivial_closure     (3)
///                           /            \
/// %b = convertToBlock %trivial_closure   nil                         (4)
///                           \            /
///                   ... uses of %optional_block ...
///                    destroy_value %optional_block
///
/// where CVT is convert_escape_to_no_escape [not_guaranteed]. We assume that
/// the %optional_block is going through a conversion sequence in PILGen meaning
/// that we should only have a single destroy of the optional block.
///
/// NOTE: There is a *lifetime gap* during the usage of the trivial_closure!
/// This means we must be careful when lifetime extending. We can only assume
/// that the underlying closure is alive immediately at the CVT. So to perform
/// our lifetime extend, we do the following:
///
/// 1. We copy and borrow optional_closure, right before the switch_enum in
/// (1).
///
/// 2. We rewrite the convert_escape_to_no_escape guaranteed to use the copy
/// instead.
///
/// 3. To make sure that even after ossa is complete, we do not move any
/// destroys above the convert_escape_to_no_escape by putting a mark_dependence
/// on %closure
///
/// 4. We insert an end_borrow, destroy for the copy at the destroy of the
/// optional block.
static bool trySwitchEnumPeephole(ConvertEscapeToNoEscapeInst *cvt) {
   auto *blockArg = dyn_cast<PILArgument>(cvt->getOperand());
   if (!blockArg)
      return false;
   auto *convertSuccessorBlock = cvt->getParent()->getSingleSuccessorBlock();
   if (!convertSuccessorBlock)
      return false;
   auto *predBB = cvt->getParent()->getSinglePredecessorBlock();
   if (!predBB)
      return false;
   auto *switchEnum1 = dyn_cast<SwitchEnumInst>(predBB->getTerminator());
   if (!switchEnum1)
      return false;
   auto *diamondSucc = getOptionalDiamondSuccessor(switchEnum1);
   if (!diamondSucc)
      return false;
   auto *switchEnum2 = dyn_cast<SwitchEnumInst>(diamondSucc->getTerminator());
   if (!switchEnum2)
      return false;
   auto *diamondSucc2 = getOptionalDiamondSuccessor(switchEnum2);
   if (!diamondSucc2)
      return false;
   if (diamondSucc2->getNumArguments() != 1)
      return false;

   // Look for the last and only destroy of the diamond succ 2's argument. This
   // is going to be the place where we destroy the lifetime extending copy.
   PILInstruction *onlyDestroy = [&]() -> PILInstruction * {
      PILInstruction *lastDestroy = nullptr;
      for (auto *use : diamondSucc2->getArgument(0)->getUses()) {
         PILInstruction *user = use->getUser();
         if (isa<ReleaseValueInst>(user) || isa<StrongReleaseInst>(user) ||
             isa<DestroyValueInst>(user)) {
            if (lastDestroy)
               return nullptr;
            lastDestroy = user;
         }
      }
      return lastDestroy;
   }();
   if (!onlyDestroy)
      return false;

   // Extend the lifetime.
   auto loc = RegularLocation::getAutoGeneratedLocation();
   PILValue copy, borrow;
   std::tie(copy, borrow) = ([&]() -> std::pair<PILValue, PILValue> {
      PILBuilderWithScope builder(switchEnum1);
      auto copy = builder.emitCopyValueOperation(loc, switchEnum1->getOperand());
      auto borrow = builder.emitBeginBorrowOperation(loc, copy);
      return {copy, borrow};
   })(); // end std::tie(copy, borrow).

   {
      PILBuilderWithScope builder(cvt);
      auto value = builder.emitExtractOptionalPayloadOperation(loc, borrow);
      cvt->setOperand(value);
      cvt->setLifetimeGuaranteed();
   }

   {
      PILBuilderWithScope builder(onlyDestroy);
      builder.emitEndBorrowOperation(loc, borrow);
      builder.emitDestroyValueOperation(loc, copy);
   }

   return true;
}

/// Look for a single destroy user and possibly unowned apply uses.
static PILInstruction *getOnlyDestroy(CopyBlockWithoutEscapingInst *cb) {
   PILInstruction *onlyDestroy = nullptr;

   for (auto *use : getNonDebugUses(cb)) {
      PILInstruction *inst = use->getUser();

      // If this an apply use, only handle unowned parameters.
      if (auto apply = FullApplySite::isa(inst)) {
         PILArgumentConvention conv = apply.getArgumentConvention(*use);
         if (conv != PILArgumentConvention::Direct_Unowned)
            return nullptr;
         continue;
      }

      // We have already seen one destroy.
      if (onlyDestroy)
         return nullptr;

      if (isa<DestroyValueInst>(inst) || isa<ReleaseValueInst>(inst) ||
          isa<StrongReleaseInst>(inst)) {
         onlyDestroy = inst;
         continue;
      }

      // Some other instruction.
      return nullptr;
   }

   if (!onlyDestroy)
      return nullptr;

   // Now look at whether the dealloc_stack or the destroy postdominates and
   // return the post dominator.
   auto *blockInit = dyn_cast<InitBlockStorageHeaderInst>(cb->getBlock());
   if (!blockInit)
      return nullptr;

   auto *asi = dyn_cast<AllocStackInst>(blockInit->getBlockStorage());
   if (!asi)
      return nullptr;
   auto *dealloc = asi->getSingleDeallocStack();
   if (!dealloc || dealloc->getParent() != onlyDestroy->getParent())
      return nullptr;

   // Return the later instruction.
   for (auto it = PILBasicBlock::iterator(onlyDestroy),
           ie = dealloc->getParent()->end();
        it != ie; ++it) {
      if (&*it == dealloc)
         return dealloc;
   }
   return onlyDestroy;
}

/// Lower a copy_block_without_escaping instruction.
///
///    This involves replacing:
///
///      %copy = copy_block_without_escaping %block withoutEscaping %closure
///
///      ...
///      destroy_value %copy
///
///    by (roughly) the instruction sequence:
///
///      %copy = copy_block %block
///
///      ...
///      destroy_value %copy
///      %e = is_escaping %closure
///      cond_fail %e
///      destroy_value %closure
static bool fixupCopyBlockWithoutEscaping(CopyBlockWithoutEscapingInst *cb,
                                          bool &modifiedCFG) {
   // Find the end of the lifetime of the copy_block_without_escaping
   // instruction.
   auto &fn = *cb->getFunction();

   // If we find a single destroy, this destroy is going to be a destroy that may
   // be in the same block as CB. It is important that we make sure that the
   // destroy is in a different block than CB or any terminating blocks to ensure
   // that we can use the SSAUpdater if needed.
   auto *singleDestroy = getOnlyDestroy(cb);
   if (singleDestroy && singleDestroy->getParent() == cb->getParent()) {
      modifiedCFG = true;
      {
         PILBuilderWithScope b(singleDestroy);
         splitBasicBlockAndBranch(b, singleDestroy, nullptr, nullptr);
      }

      {
         PILBuilderWithScope b(singleDestroy);
         auto *term = singleDestroy->getParent()->getTerminator();
         if (term->isFunctionExiting()) {
            splitBasicBlockAndBranch(b, &*std::next(singleDestroy->getIterator()),
                                     nullptr, nullptr);
         }
      }
   }

   auto sentinelClosure = cb->getClosure();
   auto loc = cb->getLoc();

   // At this point, we transform our copy_block_without_escaping into a
   // copy_block. This has a few important implications:
   //
   // 1. copy_block_without_escaping takes the sentinel value at +1. We will need
   //    to balance that +1.
   // 2. The destroy_value associated with the copy_block_without_escaping will
   //    be on the copy_block value.
   PILBuilderWithScope b(cb);
   auto *newCB = b.createCopyBlock(loc, cb->getBlock());
   cb->replaceAllUsesWith(newCB);
   cb->eraseFromParent();

   auto autoGenLoc = RegularLocation::getAutoGeneratedLocation();

   // If CB is in the entry block, we know that our definition of SentinelClosure
   // must be as well. Thus we know that we do not need to worry about loops or
   // dominance issues and can just insert destroy_values for the sentinel at the
   // lifetime end points.
   if (newCB->getParent() == newCB->getFunction()->getEntryBlock()) {
      // Our single destroy must not be in the entry block since if so, we would
      // have inserted an edge to appease the SSA updater.
      if (singleDestroy) {
         PILBuilderWithScope b(std::next(singleDestroy->getIterator()));
         PILValue v = sentinelClosure;
         PILValue isEscaping = b.createIsEscapingClosure(
            loc, v, IsEscapingClosureInst::ObjCEscaping);
         b.createCondFail(loc, isEscaping, "non-escaping closure has escaped");
         b.createDestroyValue(loc, v);
         return true;
      }

      // If we couldn't find a specific destroy_value, lifetime extend to the end
      // of the function.
      SmallVector<PILBasicBlock *, 4> ExitingBlocks;
      fn.findExitingBlocks(ExitingBlocks);
      for (auto *Block : ExitingBlocks) {
         PILBuilderWithScope B(Block->getTerminator());
         PILValue V = sentinelClosure;
         PILValue isEscaping = B.createIsEscapingClosure(
            loc, V, IsEscapingClosureInst::ObjCEscaping);
         B.createCondFail(loc, isEscaping, "non-escaping closure has escaped");
         B.createDestroyValue(loc, V);
      }

      return true;
   }

   // Otherwise, we need to be more careful since we can have loops and may not
   // transitively dominate all uses of the closure. So we:
   //
   // 1. Create an Optional<T>.none at the entry.
   // 2. Create a destroy_value(val), val = Optional<T>.some(sentinel) in the cvt
   // block.
   // 3. Create a destroy_value at all exits of the value.
   //
   // and then use the SSAUpdater to ensure that we handle loops correctly.
   auto optionalEscapingClosureTy =
      PILType::getOptionalType(sentinelClosure->getType());

   SmallVector<PILPhiArgument *, 8> insertedPhis;
   PILSSAUpdater updater(&insertedPhis);
   updater.Initialize(optionalEscapingClosureTy);

   // Create the Optional.none as the beginning available value.
   PILValue entryBlockOptionalNone;
   {
      PILBuilderWithScope b(fn.getEntryBlock()->begin());
      entryBlockOptionalNone =
         b.createOptionalNone(autoGenLoc, optionalEscapingClosureTy);
      updater.AddAvailableValue(fn.getEntryBlock(), entryBlockOptionalNone);
   }
   assert(entryBlockOptionalNone);

   // Then create the Optional.some(closure sentinel).
   //
   // NOTE: We return the appropriate insertion point to insert the destroy_value
   // before the value (to ensure we handle loops). We need to get all available
   // values first though.
   auto *initialValue = [&]() -> EnumInst * {
      PILBuilderWithScope b(newCB);
      // Create the closure sentinel (the copy_block_without_escaping closure
      // operand consumed at +1, so we don't need a copy) to it.
      auto *result = b.createOptionalSome(autoGenLoc, sentinelClosure,
                                          optionalEscapingClosureTy);
      updater.AddAvailableValue(result->getParent(), result);
      return result;
   }();

   // If we had a single destroy, creating a .none after it and add that as a
   // value to the SSA updater.
   if (singleDestroy) {
      PILBuilderWithScope b(std::next(singleDestroy->getIterator()));
      auto *result = b.createOptionalNone(autoGenLoc, optionalEscapingClosureTy);
      updater.AddAvailableValue(result->getParent(), result);
   }

   // Now that we have all of our available values, insert a destroy_value before
   // the initial Optional.some value using the SSA updater to ensure that we
   // handle loops correctly.
   {
      PILValue v = updater.GetValueInMiddleOfBlock(initialValue->getParent());
      PILBuilderWithScope(initialValue).createDestroyValue(autoGenLoc, v);
   }

   // And insert an is_escaping_closure, cond_fail, destroy_value at each of the
   // lifetime end points. This ensures we do not expand our lifetime too much.
   if (singleDestroy) {
      PILBuilderWithScope b(std::next(singleDestroy->getIterator()));
      PILValue v = updater.GetValueInMiddleOfBlock(singleDestroy->getParent());
      PILValue isEscaping =
         b.createIsEscapingClosure(loc, v, IsEscapingClosureInst::ObjCEscaping);
      b.createCondFail(loc, isEscaping, "non-escaping closure has escaped");
      b.createDestroyValue(loc, v);
   }

   // Then to be careful with regards to loops, insert at each of the destroy
   // blocks destroy_value to ensure that we obey ownership invariants.
   {
      SmallVector<PILBasicBlock *, 4> exitingBlocks;
      findReachableExitBlocks(newCB, exitingBlocks);

      for (auto *block : exitingBlocks) {
         auto *safeDestructionPt = getDeinitSafeClosureDestructionPoint(block);
         PILValue v = updater.GetValueAtEndOfBlock(block);
         PILBuilderWithScope(safeDestructionPt).createDestroyValue(autoGenLoc, v);
      }
   }

   // Finally, we need to prune phis inserted by the SSA updater that only take
   // the .none from the entry block.
   //
   // TODO: Should we sort inserted phis before or after we initialize
   // the worklist or maybe backwards? We should investigate how the
   // SSA updater adds phi nodes to this list to resolve this question.
   cleanupDeadTrivialPhiArgs(entryBlockOptionalNone, insertedPhis);

   return true;
}

static bool fixupClosureLifetimes(PILFunction &fn, bool &checkStackNesting,
                                  bool &modifiedCFG) {
   bool changed = false;

   // tryExtendLifetimeToLastUse uses a cache of recursive instruction use
   // queries.
   llvm::DenseMap<PILInstruction *, PILInstruction *> memoizedQueries;

   for (auto &block : fn) {
      auto i = block.begin();
      while (i != block.end()) {
         PILInstruction *inst = &*i;
         ++i;

         // Handle, copy_block_without_escaping instructions.
         if (auto *cb = dyn_cast<CopyBlockWithoutEscapingInst>(inst)) {
            if (fixupCopyBlockWithoutEscaping(cb, modifiedCFG)) {
               changed = true;
            }
            continue;
         }

         // Otherwise, look at convert_escape_to_noescape [not_guaranteed]
         // instructions.
         auto *cvt = dyn_cast<ConvertEscapeToNoEscapeInst>(inst);
         if (!cvt || cvt->isLifetimeGuaranteed())
            continue;

         // First try to peephole a known pattern.
         if (!DisableConvertEscapeToNoEscapeSwitchEnumPeephole) {
            if (trySwitchEnumPeephole(cvt)) {
               changed = true;
               continue;
            }
         }

         if (tryExtendLifetimeToLastUse(cvt, memoizedQueries, i)) {
            changed = true;
            checkStackNesting = true;
            continue;
         }

         // Otherwise, extend the lifetime of the operand to the end of the
         // function.
         extendLifetimeToEndOfFunction(fn, cvt);
         changed = true;
      }
   }
   return changed;
}

/// Fix-up the lifetime of the escaping closure argument of
/// convert_escape_to_noescape [not_guaranteed] instructions.
///
/// convert_escape_to_noescape [not_guaranteed] assume that someone guarantees
/// the lifetime of the operand for the duration of the trivial closure result.
/// PILGen does not guarantee this for '[not_guaranteed]' instructions so we
/// ensure it here.
namespace {
class ClosureLifetimeFixup : public PILFunctionTransform {

   /// The entry point to the transformation.
   void run() override {
      // Don't rerun diagnostics on deserialized functions.
      if (getFunction()->wasDeserializedCanonical())
         return;

      // Fixup convert_escape_to_noescape [not_guaranteed] and
      // copy_block_without_escaping instructions.

      bool checkStackNesting = false;
      bool modifiedCFG = false;

      if (fixupClosureLifetimes(*getFunction(), checkStackNesting, modifiedCFG)) {
         if (checkStackNesting){
            StackNesting sn;
            modifiedCFG =
               sn.correctStackNesting(getFunction()) == StackNesting::Changes::CFG;
         }
         if (modifiedCFG)
            invalidateAnalysis(PILAnalysis::InvalidationKind::FunctionBody);
         else
            invalidateAnalysis(PILAnalysis::InvalidationKind::CallsAndInstructions);
      }
      LLVM_DEBUG(getFunction()->verify());

   }

};
} // end anonymous namespace

PILTransform *polar::createClosureLifetimeFixup() {
   return new ClosureLifetimeFixup();
}
